package TestCases;

import com.clickhouse.jdbc.ClickHouseConnection;
import com.clickhouse.jdbc.ClickHouseDataSource;
import com.clickhouse.jdbc.ClickHousePreparedStatement;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.IOException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.sql.DriverManager;
import java.sql.SQLException;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;

public class ClickHouseBatchInsertUtil {
    static Logger log = LoggerUtils.getLogger(ClickHouseBatchInsertUtil.class);
    static String CLICKHOUSE_URL = ConfigUtil.getProperty("ClickHouse.jdbc.url", ConstantsUtil.CONFIG_JDBC);
    static String CLICKHOUSE_USERNAME = ConfigUtil.getProperty("ClickHouse.jdbc.userName", ConstantsUtil.CONFIG_JDBC);
    static String CLICKHOUSE_PASSWORD = ConfigUtil.getProperty("ClickHouse.jdbc.passWord", ConstantsUtil.CONFIG_JDBC);
    static String CLICKHOUSE_TABLE_NAME1 = ConfigUtil.getProperty("ClickHouse.jdbc.tableName1", ConstantsUtil.CONFIG_JDBC);
    static int CLICKHOUSE_BATCH_SIZE = Integer.parseInt(ConfigUtil.getProperty("ClickHouse.jdbc.batchSize", ConstantsUtil.CONFIG_JDBC));


    static String insertSQL1 = ConfigUtil.getProperty("insertSQL1", ConstantsUtil.CONFIG_COMMON);
    static String insertSQL2 = ConfigUtil.getProperty("insertSQL2", ConstantsUtil.CONFIG_COMMON);


    public static ClickHouseConnection connection() throws SQLException {
        try {
            ClickHouseConnection connection = (ClickHouseConnection) DriverManager.getConnection(CLICKHOUSE_URL, CLICKHOUSE_USERNAME, CLICKHOUSE_PASSWORD);
            if (connection != null) {
                log.info("Successfully to make connection!");
                return connection;
            } else {
                log.info("Failed to make connection!");
            }
        } catch (SQLException e) {
            System.err.println("Error while connecting to ClickHouse: " + e.getMessage());
        }
        return null;
    }

    public static void close(ClickHouseConnection connection, ClickHousePreparedStatement pstmt) throws SQLException {
        if (connection != null && !connection.isClosed()) {
            connection.close();
        }
        if (pstmt != null && !pstmt.isClosed()) {
            pstmt.close();
        }
        log.info("Successfully close to ClickHouse!");
    }

    public static void test() throws SQLException {
        String query = "SELECT * FROM audit_record";

        try (Connection connection = DriverManager.getConnection(CLICKHOUSE_URL, CLICKHOUSE_USERNAME, CLICKHOUSE_PASSWORD); Statement statement = connection.createStatement()) {
            ResultSet resultSet = statement.executeQuery(query);
            while (resultSet.next()) {
                Long id = resultSet.getLong("id");
                String tenantId = resultSet.getString("tenantId");
                log.info("id: " + id + ", tenantId: " + tenantId);
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    public static void batchInsertJson1(ClickHouseConnection connection, ClickHousePreparedStatement pstmt, String jsonLog, int batchSize) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            AuditRecord audit_record = objectMapper.readValue(jsonLog, AuditRecord.class);
            Object[] rowData = {audit_record.getId(), audit_record.getTenantId(), audit_record.getHappenTime(), // 需要将日期字符串转换为DateTime对象
                    audit_record.getSrcIp(), audit_record.getSrcPort(), audit_record.getSrcMac(), audit_record.getSystemUser(), audit_record.getSystemHost(), audit_record.getDestIp(), audit_record.getDestPort(), audit_record.getDestMac(), audit_record.getVisitTool(), audit_record.getAppAccount(), audit_record.getSessionId(), audit_record.getDbType(), audit_record.getDbUser(), audit_record.getOperType(), audit_record.getDbName(), audit_record.getTableName(), audit_record.getFieldName(), audit_record.getSqlTemplateId(), audit_record.getOperSentence(), audit_record.getOperSentenceLen(), audit_record.getSqlBindValue(), audit_record.getRowNum(), audit_record.getSqlExecTime(), audit_record.getSqlResponse(), audit_record.getReturnContent(), audit_record.getReturnContentLen(), audit_record.getDealState(), audit_record.getProtectObjectName(), audit_record.getRuleType(), audit_record.getRuleName(), audit_record.getRiskLev(), audit_record.getHwId(), audit_record.getRelatedInfo(), audit_record.getApprovalStatus()};
            pstmt.clearParameters(); // 清除已存在的参数
            for (int i = 0; i < rowData.length; i++) {
                pstmt.setObject(i + 1, rowData[i]);
            }
            pstmt.addBatch();
            if (batchSize == CLICKHOUSE_BATCH_SIZE) {
                log.info("当前已添加批量数据：" + batchSize + "=>到达了批处理大小阈值" + CLICKHOUSE_BATCH_SIZE + "，开始执行批量提交！");
                pstmt.executeBatch();
//                    pstmt.clearBatch(); // 清除已添加的批量操作
            }
        } catch (SQLException | IOException e) {
            // 记录错误信息
            log.error("Error during batch insert: " + e.getMessage());
            // 深度处理错误，比如重试、跳过、记录详细日志等
        } catch (Exception e) {
            // 记录解析或其他异常
            log.error("Error parsing JSON or other exception: " + e.getMessage());
        } finally {
            try {
                if (pstmt != null && !pstmt.isClosed()) {
                    pstmt.executeBatch(); // 如果在循环结束时仍有未执行的批处理，此处将会执行
                    pstmt.clearBatch();
                }
            } catch (SQLException e) {
                System.err.println("Error closing resources: " + e.getMessage());
            }
        }
    }

    public static void batchInsertJson2(ClickHouseConnection connection, ClickHousePreparedStatement pstmt, List<String> jsonLogs) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            for (String jsonLog : jsonLogs) {
                AuditRecord audit_record = objectMapper.readValue(jsonLog, AuditRecord.class);
                Object[] rowData = {audit_record.getId(), audit_record.getTenantId(), audit_record.getHappenTime(), // 需要将日期字符串转换为DateTime对象
                        audit_record.getSrcIp(), audit_record.getSrcPort(), audit_record.getSrcMac(), audit_record.getSystemUser(), audit_record.getSystemHost(), audit_record.getDestIp(), audit_record.getDestPort(), audit_record.getDestMac(), audit_record.getVisitTool(), audit_record.getAppAccount(), audit_record.getSessionId(), audit_record.getDbType(), audit_record.getDbUser(), audit_record.getOperType(), audit_record.getDbName(), audit_record.getTableName(), audit_record.getFieldName(), audit_record.getSqlTemplateId(), audit_record.getOperSentence(), audit_record.getOperSentenceLen(), audit_record.getSqlBindValue(), audit_record.getRowNum(), audit_record.getSqlExecTime(), audit_record.getSqlResponse(), audit_record.getReturnContent(), audit_record.getReturnContentLen(), audit_record.getDealState(), audit_record.getProtectObjectName(), audit_record.getRuleType(), audit_record.getRuleName(), audit_record.getRiskLev(), audit_record.getHwId(), audit_record.getRelatedInfo(), audit_record.getApprovalStatus()};
                for (int i = 0; i < rowData.length; i++) {
                    pstmt.setObject(i + 1, rowData[i]);
                }
                pstmt.addBatch();
            }
            pstmt.executeBatch();
        } catch (SQLException | IOException e) {
            // 记录错误信息
            log.error("Error during batch insert: " + e.getMessage());
            // 深度处理错误，比如重试、跳过、记录详细日志等
        } catch (Exception e) {
            // 记录解析或其他异常
            log.error("Error parsing JSON or other exception: " + e.getMessage());
        }
    }

    public static void batchInsertJson3(int for_riskLev, List<String> jsonLogs, int ThreadsNums) {
        int numThreads = ThreadsNums;
        log.info("Start the thread pool");
        ExecutorService executor = Executors.newFixedThreadPool(numThreads);
        int totalSize = jsonLogs.size();
        int chunkSize = (totalSize + numThreads - 1) / numThreads; // Divide logs evenly among threads
        for (int i = 0; i < totalSize; i += chunkSize) {
            int end = Math.min(i + chunkSize, totalSize);
            List<String> sublist = jsonLogs.subList(i, end);
            executor.submit(() -> {
                try (ClickHouseConnection connection = ClickHouseBatchInsertUtil.connection()) {
                    batchInsertTask(for_riskLev, connection, sublist);
                } catch (SQLException | IOException e) {
                    e.printStackTrace();
                }
            });
        }
        executor.shutdown();
        try {
            // Wait for all tasks to finish
            if (!executor.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS)) {
                log.error("Not all tasks completed.");
            }
        } catch (InterruptedException e) {
            log.error("Thread was interrupted.");
        }
    }

    private static void batchInsertTask(int for_riskLev, ClickHouseConnection connection, List<String> jsonLogs) throws IOException, SQLException {
        ObjectMapper objectMapper = new ObjectMapper();
        log.info("线程" + Thread.currentThread().getName() + "开始处理");
        ClickHousePreparedStatement pstmt1 = null;
        ClickHousePreparedStatement pstmt2 = null;
        try {
            if (for_riskLev < 3) {
                pstmt1 = (ClickHousePreparedStatement) connection.prepareStatement(insertSQL1);
                pstmt2 = (ClickHousePreparedStatement) connection.prepareStatement(insertSQL2);
            } else {
                pstmt1 = (ClickHousePreparedStatement) connection.prepareStatement(insertSQL1);
            }

            for (String jsonLog : jsonLogs) {
                AuditRecord audit_record = objectMapper.readValue(jsonLog, AuditRecord.class);
                Object[] rowData = {audit_record.getId(), audit_record.getTenantId(), audit_record.getHappenTime(),
                        audit_record.getSrcIp(), audit_record.getSrcPort(), audit_record.getSrcMac(), audit_record.getSystemUser(), audit_record.getSystemHost(), audit_record.getDestIp(), audit_record.getDestPort(), audit_record.getDestMac(), audit_record.getVisitTool(), audit_record.getAppAccount(), audit_record.getSessionId(), audit_record.getDbType(), audit_record.getDbUser(), audit_record.getOperType(), audit_record.getDbName(), audit_record.getTableName(), audit_record.getFieldName(), audit_record.getSqlTemplateId(), audit_record.getOperSentence(), audit_record.getOperSentenceLen(), audit_record.getSqlBindValue(), audit_record.getRowNum(), audit_record.getSqlExecTime(), audit_record.getSqlResponse(), audit_record.getReturnContent(), audit_record.getReturnContentLen(), audit_record.getDealState(), audit_record.getProtectObjectName(), audit_record.getRuleType(), audit_record.getRuleName(), audit_record.getRiskLev(), audit_record.getHwId(), audit_record.getRelatedInfo(), audit_record.getApprovalStatus()};

                if (for_riskLev < 3) {
                    for (int i = 0; i < rowData.length; i++) {
                        pstmt1.setObject(i + 1, rowData[i]);
                        pstmt2.setObject(i + 1, rowData[i]);
                    }
                    pstmt1.addBatch();
                    pstmt2.addBatch();
                } else {
                    for (int i = 0; i < rowData.length; i++) {
                        pstmt1.setObject(i + 1, rowData[i]);
                    }
                    pstmt1.addBatch();
                }
            }

            if (for_riskLev < 3) {
                pstmt1.executeBatch();
                pstmt2.executeBatch();
            } else {
                pstmt1.executeBatch();
            }
        } finally {
            if (pstmt1 != null) pstmt1.close();
            if (pstmt2 != null) pstmt2.close();
            log.info("线程" + Thread.currentThread().getName() + "处理完毕");
            connection.close();
        }
    }


    // 示例如何使用该工具方法
    public static void main(String[] args) throws Exception {
//    	test();
        String json = "{\r\n" + "		\"id\" : 167526723021000030,\r\n" + "		\"tenantId\" : \"0\",\r\n" + "		\"happenTime\" : \"2024-04-11 21:45:17\",\r\n" + "		\"srcIp\" : \"193.2.1.1\",\r\n" + "		\"srcPort\" : 62893,\r\n" + "		\"srcMac\" : \"2C:56:DC:79:A1:96\",\r\n" + "		\"systemUser\" : \" \",\r\n" + "		\"systemHost\" : \" \",\r\n" + "		\"destIp\" : \"172.18.1.82\",\r\n" + "		\"destPort\" : 3306,\r\n" + "		\"destMac\" : \"00:E2:69:09:D4:C1\",\r\n" + "		\"visitTool\" : \" \",\r\n" + "		\"appAccount\" : \" \",\r\n" + "		\"sessionId\" : \"8eaeca63-90d4-4055-a438-5ac9945944b8\",\r\n" + "		\"dbType\" : 2,\r\n" + "		\"dbUser\" : \"dbUser0\",\r\n" + "		\"operType\" : \"select\",\r\n" + "		\"dbName\" : \"\",\r\n" + "		\"tableName\" : \"admin_user\",\r\n" + "		\"fieldName\" : \"所有字段\",\r\n" + "		\"sqlTemplateId\" : \"7163112700047907533\",\r\n" + "		\"operSentence\" : \"select * from admin_us where fid='cvjg35m3k0coee5' and fd = 20000000002 and fd1='rabbit dragon monkey'\",\r\n" + "		\"operSentenceLen\" : 102,\r\n" + "		\"sqlBindValue\" : \" \",\r\n" + "		\"rowNum\" : 2,\r\n" + "		\"sqlExecTime\" : 0.056,\r\n" + "		\"sqlResponse\" : 1,\r\n" + "		\"returnContent\" : \"user_id @@@username @@@password @@@realname @@@email @@@mobile @@@login_count @@@last_login_ip @@@last_login_time @@@status @@@create_time $$$1@@@u0001@@@6f493fbddf9107797f5044bb229ac6ee@@@Administrator001@@@admin@mail.com@@@0001@@@48@@@42.196.163.128@@@2016-05-01 09:53:36@@@1@@@2013-12-25 15:58:34$$$8@@@u0002@@@e10adc3949ba59abbe56e057f20f883e@@@Guest02@@@000002@@@0000002@@@5664@@@42.196.163.128@@@2016-05-01 09:28:46@@@1@@@2014-03-$$$\",\r\n" + "		\"returnContentLen\" : 439,\r\n" + "		\"dealState\" : 1,\r\n" + "		\"protectObjectName\" : \"172.18.1.82\",\r\n" + "		\"ruleType\" : 0,\r\n" + "		\"ruleName\" : \"rule0\",\r\n" + "		\"riskLev\" : 2,\r\n" + "		\"hwId\" : \"\",\r\n" + "		\"relatedInfo\" : \" \",\r\n" + "		\"approvalStatus\" : 0\r\n" + "	}";
        ClickHouseConnection connection = connection();
        List<String> jsons = new ArrayList<>();
        jsons.add(json);
        jsons.add(json);
        jsons.add(json);
//        List<String> jsonLogs = Arrays.asList(jsonString);
        log.info("" + jsons);
//        batchInsertJson(jsonLogs);
        String insertQuery = "INSERT INTO " + CLICKHOUSE_TABLE_NAME1 + " (id, tenantId, happenTime, srcIp, srcPort, srcMac, systemUser, systemHost, destIp, destPort, destMac, visitTool, appAccount, sessionId, dbType, dbUser, operType, dbName, tableName, fieldName, sqlTemplateId, operSentence, operSentenceLen, sqlBindValue, rowNum, sqlExecTime, sqlResponse, returnContent, returnContentLen, dealState, protectObjectName, ruleType, ruleName, riskLev, hwId, relatedInfo, approvalStatus) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
        ClickHousePreparedStatement pstmt = (ClickHousePreparedStatement) connection.prepareStatement(insertQuery);
        try {
            batchInsertJson1(connection, pstmt, json, 2);
//        	batchInsertJson2(connection,pstmt, jsons,2);
        } finally {
            close(connection, pstmt);
        }
    }
}
